from grid_env_ideal_obs_repeat_task import *
from maze_factory import *
from replay_config import *
import argparse
import json
import time
import svgwrite
from datetime import datetime
import os

# load landscape and states from file
def load_task(pth = "./logs/task.json"):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        print("state: ", state)
        print("goal: ", goal)
        print("landscape: ", landscape)
    return landscape, state, goal

# save current landscape as json file
def save_current_task(landscape, start_x, start_y, goal_x, goal_y, pth = "./logs/landscape.json"):
    landscape_ = []
    for j in range(landscape[0].shape[0]):
        landscape_.append(int(landscape[0][j]))

    with open(pth, "w") as f:
        json.dump({"data": landscape_, 
                   "state": [start_x, start_y],
                   "goal": [goal_x, goal_y]}, f)


def render_single_patch(obs_patch, grid_size_display=40):
    """绘制单个3x3观测patch，使用与完整地图相同的风格"""
    patch_size = 3
    img = np.zeros((patch_size * grid_size_display, patch_size * grid_size_display, 3), np.uint8)
    
    for j in range(patch_size):
        for i in range(patch_size):
            obs_value = obs_patch[j * patch_size + i]
            
            if obs_value == 0:  # 自由空间
                # 自由空间 - 深蓝灰色（与完整地图一致）
                cv2.rectangle(img, (i * grid_size_display, j * grid_size_display), 
                            (i * grid_size_display + grid_size_display, j * grid_size_display + grid_size_display), 
                            (30, 40, 50), -1)
                
                # 立体边框效果：深色调
                cv2.line(img, (i * grid_size_display, j * grid_size_display), 
                       ((i + 1) * grid_size_display - 1, j * grid_size_display), 
                       (50, 70, 90), 4)
                cv2.line(img, (i * grid_size_display, j * grid_size_display), 
                       (i * grid_size_display, (j + 1) * grid_size_display - 1), 
                       (50, 70, 90), 4)
                cv2.line(img, (i * grid_size_display, (j + 1) * grid_size_display - 1), 
                       ((i + 1) * grid_size_display - 1, (j + 1) * grid_size_display - 1), 
                       (15, 25, 35), 4)
                cv2.line(img, ((i + 1) * grid_size_display - 1, j * grid_size_display), 
                       ((i + 1) * grid_size_display - 1, (j + 1) * grid_size_display - 1), 
                       (15, 25, 35), 4)
            elif obs_value == 1:  # 障碍物
                # 障碍物 - 浅灰色（与完整地图一致）
                cv2.rectangle(img, (i * grid_size_display, j * grid_size_display), 
                            (i * grid_size_display + grid_size_display, j * grid_size_display + grid_size_display), 
                            (200, 200, 210), -1)
                
                # 石墙立体边框效果
                cv2.line(img, (i * grid_size_display, j * grid_size_display), 
                       ((i + 1) * grid_size_display - 1, j * grid_size_display), 
                       (230, 230, 240), 4)
                cv2.line(img, (i * grid_size_display, j * grid_size_display), 
                       (i * grid_size_display, (j + 1) * grid_size_display - 1), 
                       (230, 230, 240), 4)
                cv2.line(img, (i * grid_size_display, (j + 1) * grid_size_display - 1), 
                       ((i + 1) * grid_size_display - 1, (j + 1) * grid_size_display - 1), 
                       (170, 170, 180), 4)
                cv2.line(img, ((i + 1) * grid_size_display - 1, j * grid_size_display), 
                       ((i + 1) * grid_size_display - 1, (j + 1) * grid_size_display - 1), 
                       (170, 170, 180), 4)
                
                # 添加浅色纹理效果
                cv2.circle(img, (i * grid_size_display + 12, j * grid_size_display + 12), 2, (180, 180, 190), -1)
                cv2.circle(img, (i * grid_size_display + 28, j * grid_size_display + 28), 2, (180, 180, 190), -1)
            elif obs_value == -2:  # 目标位置
                # 目标位置特殊标记
                cv2.rectangle(img, (i * grid_size_display, j * grid_size_display), 
                            (i * grid_size_display + grid_size_display, j * grid_size_display + grid_size_display), 
                            (0, 100, 0), -1)
                cv2.line(img, (i * grid_size_display, j * grid_size_display), 
                       ((i + 1) * grid_size_display - 1, j * grid_size_display), 
                       (0, 150, 0), 4)
                cv2.line(img, (i * grid_size_display, j * grid_size_display), 
                       (i * grid_size_display, (j + 1) * grid_size_display - 1), 
                       (0, 150, 0), 4)
                cv2.line(img, (i * grid_size_display, (j + 1) * grid_size_display - 1), 
                       ((i + 1) * grid_size_display - 1, (j + 1) * grid_size_display - 1), 
                       (0, 120, 0), 4)
                cv2.line(img, ((i + 1) * grid_size_display - 1, j * grid_size_display), 
                       ((i + 1) * grid_size_display - 1, (j + 1) * grid_size_display - 1), 
                       (0, 120, 0), 4)
    
    return img

def create_patches_montage(observation_patches, separator_height=5):
    """创建observation patches的纵向拼接图"""
    if not observation_patches:
        return None
    
    grid_size_display = 40
    patch_size = 3
    patch_height = patch_size * grid_size_display
    patch_width = patch_size * grid_size_display
    
    # 计算总高度：所有patch + 分隔符
    total_height = len(observation_patches) * patch_height + (len(observation_patches) - 1) * separator_height
    total_width = patch_width
    
    # 创建拼接图画布
    montage_img = np.zeros((total_height, total_width, 3), np.uint8)
    
    # 逐个绘制每个patch
    current_y = 0
    for i, obs_patch in enumerate(observation_patches):
        # 绘制单个patch
        patch_img = render_single_patch(obs_patch, grid_size_display)
        
        # 将patch复制到拼接图中
        montage_img[current_y:current_y + patch_height, 0:patch_width] = patch_img
        
        # 更新位置（添加分隔符）
        current_y += patch_height + separator_height
    
    return montage_img

def save_patches_montage_svg(observation_patches, pth, separator_height=5):
    """保存observation patches拼接图为SVG格式"""
    if not observation_patches:
        return
    
    grid_size_display = 40
    patch_size = 3
    patch_height = patch_size * grid_size_display
    patch_width = patch_size * grid_size_display
    
    # 计算总高度
    total_height = len(observation_patches) * patch_height + (len(observation_patches) - 1) * separator_height
    total_width = patch_width
    
    # 创建SVG画布
    dwg = svgwrite.Drawing(pth, size=(total_width, total_height))
    
    # 逐个绘制每个patch
    current_y = 0
    for patch_idx, obs_patch in enumerate(observation_patches):
        for j in range(patch_size):
            for i in range(patch_size):
                x = i * grid_size_display
                y = current_y + j * grid_size_display
                obs_value = obs_patch[j * patch_size + i]
                
                if obs_value == 0:  # 自由空间
                    dwg.add(dwg.rect(insert=(x, y), size=(grid_size_display, grid_size_display), 
                                   fill='rgb(30,40,50)', stroke='none'))
                    # 立体边框效果
                    dwg.add(dwg.line(start=(x, y), end=(x+grid_size_display-1, y), 
                                   stroke='rgb(50,70,90)', stroke_width=4))
                    dwg.add(dwg.line(start=(x, y), end=(x, y+grid_size_display-1), 
                                   stroke='rgb(50,70,90)', stroke_width=4))
                    dwg.add(dwg.line(start=(x, y+grid_size_display-1), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(15,25,35)', stroke_width=4))
                    dwg.add(dwg.line(start=(x+grid_size_display-1, y), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(15,25,35)', stroke_width=4))
                elif obs_value == 1:  # 障碍物
                    dwg.add(dwg.rect(insert=(x, y), size=(grid_size_display, grid_size_display), 
                                   fill='rgb(200,200,210)', stroke='none'))
                    # 立体边框效果
                    dwg.add(dwg.line(start=(x, y), end=(x+grid_size_display-1, y), 
                                   stroke='rgb(230,230,240)', stroke_width=4))
                    dwg.add(dwg.line(start=(x, y), end=(x, y+grid_size_display-1), 
                                   stroke='rgb(230,230,240)', stroke_width=4))
                    dwg.add(dwg.line(start=(x, y+grid_size_display-1), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(170,170,180)', stroke_width=4))
                    dwg.add(dwg.line(start=(x+grid_size_display-1, y), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(170,170,180)', stroke_width=4))
                    # 添加浅色纹理点
                    dwg.add(dwg.circle(center=(x+12, y+12), r=2, fill='rgb(180,180,190)'))
                    dwg.add(dwg.circle(center=(x+28, y+28), r=2, fill='rgb(180,180,190)'))
                elif obs_value == -2:  # 目标位置
                    dwg.add(dwg.rect(insert=(x, y), size=(grid_size_display, grid_size_display), 
                                   fill='rgb(0,100,0)', stroke='none'))
                    dwg.add(dwg.line(start=(x, y), end=(x+grid_size_display-1, y), 
                                   stroke='rgb(0,150,0)', stroke_width=4))
                    dwg.add(dwg.line(start=(x, y), end=(x, y+grid_size_display-1), 
                                   stroke='rgb(0,150,0)', stroke_width=4))
                    dwg.add(dwg.line(start=(x, y+grid_size_display-1), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(0,120,0)', stroke_width=4))
                    dwg.add(dwg.line(start=(x+grid_size_display-1, y), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(0,120,0)', stroke_width=4))
        
        # 更新Y位置（添加分隔符）
        current_y += patch_height + separator_height
    
    # 保存SVG文件
    dwg.save()
    print(f"Patches montage SVG saved to {pth}")

def save_svg_snapshot(grid, trajectory, init_pos, goal_pos, agent_pos, agent_action, perturbation_arrows, pth, path_color="green"):
    """将当前画面保存为高质量SVG矢量图"""
    grid_size_display = 40
    width, height = grid.shape[0], grid.shape[1]
    svg_width = width * grid_size_display
    svg_height = height * grid_size_display
    
    # 创建SVG画布
    dwg = svgwrite.Drawing(pth, size=(svg_width, svg_height))
    
    # 计算agent位置用于可视范围判断
    agent_grid_x = int((agent_pos[0] - 20) // 40)
    agent_grid_y = int((agent_pos[1] - 20) // 40)
    
    # 绘制网格 - 加上立体效果和可视范围区分
    for j in range(width):
        for i in range(height):
            x = i * grid_size_display
            y = j * grid_size_display
            
            # 判断是否在可视范围内
            in_vision = (abs(j - agent_grid_x) <= 1 and abs(i - agent_grid_y) <= 1)
            
            if grid[j,i] == 1:
                if in_vision:
                    # 可视范围内的自由空间 - 改为深色（与OpenCV一致）
                    dwg.add(dwg.rect(insert=(x, y), size=(grid_size_display, grid_size_display), 
                                   fill='rgb(40,50,40)', stroke='none'))
                    # 立体边框效果：深色调
                    dwg.add(dwg.line(start=(x, y), end=(x+grid_size_display-1, y), 
                                   stroke='rgb(70,90,70)', stroke_width=4))  # 上边
                    dwg.add(dwg.line(start=(x, y), end=(x, y+grid_size_display-1), 
                                   stroke='rgb(70,90,70)', stroke_width=4))  # 左边
                    dwg.add(dwg.line(start=(x, y+grid_size_display-1), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(20,30,20)', stroke_width=4))  # 下边
                    dwg.add(dwg.line(start=(x+grid_size_display-1, y), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(20,30,20)', stroke_width=4))  # 右边
                else:
                    # 不可视范围内的自由空间 - 深蓝灰色（与OpenCV一致）
                    dwg.add(dwg.rect(insert=(x, y), size=(grid_size_display, grid_size_display), 
                                   fill='rgb(30,40,50)', stroke='none'))
                    # 立体边框效果
                    dwg.add(dwg.line(start=(x, y), end=(x+grid_size_display-1, y), 
                                   stroke='rgb(50,70,90)', stroke_width=4))  # 上边
                    dwg.add(dwg.line(start=(x, y), end=(x, y+grid_size_display-1), 
                                   stroke='rgb(50,70,90)', stroke_width=4))  # 左边
                    dwg.add(dwg.line(start=(x, y+grid_size_display-1), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(15,25,35)', stroke_width=4))  # 下边
                    dwg.add(dwg.line(start=(x+grid_size_display-1, y), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(15,25,35)', stroke_width=4))  # 右边
            else:
                if in_vision:
                    # 可视范围内的障碍物 - 改为浅米色（与OpenCV一致）
                    dwg.add(dwg.rect(insert=(x, y), size=(grid_size_display, grid_size_display), 
                                   fill='rgb(220,220,200)', stroke='none'))
                    # 立体边框效果
                    dwg.add(dwg.line(start=(x, y), end=(x+grid_size_display-1, y), 
                                   stroke='rgb(250,250,230)', stroke_width=4))  # 上边
                    dwg.add(dwg.line(start=(x, y), end=(x, y+grid_size_display-1), 
                                   stroke='rgb(250,250,230)', stroke_width=4))  # 左边
                    dwg.add(dwg.line(start=(x, y+grid_size_display-1), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(190,190,170)', stroke_width=4))  # 下边
                    dwg.add(dwg.line(start=(x+grid_size_display-1, y), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(190,190,170)', stroke_width=4))  # 右边
                    # 添加浅色纹理点
                    dwg.add(dwg.circle(center=(x+12, y+12), r=2, fill='rgb(180,180,190)'))
                    dwg.add(dwg.circle(center=(x+28, y+28), r=2, fill='rgb(180,180,190)'))
                else:
                    # 不可视范围内的障碍物 - 改为浅灰色（与OpenCV一致）
                    dwg.add(dwg.rect(insert=(x, y), size=(grid_size_display, grid_size_display), 
                                   fill='rgb(200,200,210)', stroke='none'))
                    # 立体边框效果
                    dwg.add(dwg.line(start=(x, y), end=(x+grid_size_display-1, y), 
                                   stroke='rgb(230,230,240)', stroke_width=4))  # 上边
                    dwg.add(dwg.line(start=(x, y), end=(x, y+grid_size_display-1), 
                                   stroke='rgb(230,230,240)', stroke_width=4))  # 左边
                    dwg.add(dwg.line(start=(x, y+grid_size_display-1), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(170,170,180)', stroke_width=4))  # 下边
                    dwg.add(dwg.line(start=(x+grid_size_display-1, y), end=(x+grid_size_display-1, y+grid_size_display-1), 
                                   stroke='rgb(170,170,180)', stroke_width=4))  # 右边
                    # 添加浅色纹理点
                    dwg.add(dwg.circle(center=(x+12, y+12), r=2, fill='rgb(180,180,190)'))
                    dwg.add(dwg.circle(center=(x+28, y+28), r=2, fill='rgb(180,180,190)'))
    
    # 绘制轨迹路径 - 支持颜色选择，与OpenCV界面一致
    if len(trajectory) > 1:
        complete_trajectory = [init_pos] + trajectory
        
        # 绘制路径阴影
        for i in range(len(complete_trajectory)-1):
            dwg.add(dwg.line(start=(complete_trajectory[i][1]+2, complete_trajectory[i][0]+2), 
                           end=(complete_trajectory[i+1][1]+2, complete_trajectory[i+1][0]+2),
                           stroke='rgb(40,40,40)', stroke_width=8))
        
        # 绘制路径主线条 - 根据path_color选择颜色方案
        for i in range(len(complete_trajectory)-1):
            progress = i / max(1, len(complete_trajectory) - 2)
            
            if path_color == "green":
                # 绿色方案 (SVG RGB格式)
                color_r = int(0 + progress * 100)
                color_g = int(100 + progress * 155)
                color_b = int(0 + progress * 100)
                center_line_color = 'rgb(150,255,150)'
            else:  # red
                # 红色方案 (SVG RGB格式)
                color_r = int(100 + progress * 155)
                color_g = int(0 + progress * 100)
                color_b = int(0 + progress * 100)
                center_line_color = 'rgb(255,150,150)'
            
            dwg.add(dwg.line(start=(complete_trajectory[i][1], complete_trajectory[i][0]), 
                           end=(complete_trajectory[i+1][1], complete_trajectory[i+1][0]),
                           stroke=f'rgb({color_r},{color_g},{color_b})', stroke_width=6))
        
        # 绘制路径中心的亮线
        for i in range(len(complete_trajectory)-1):
            dwg.add(dwg.line(start=(complete_trajectory[i][1], complete_trajectory[i][0]), 
                           end=(complete_trajectory[i+1][1], complete_trajectory[i+1][0]),
                           stroke=center_line_color, stroke_width=2))
        
        # 绘制路径点
        for i, point in enumerate(complete_trajectory):
            progress = i / max(1, len(complete_trajectory) - 1)
            
            if path_color == "green":
                dot_r = int(50 + progress * 100)
                dot_g = int(150 + progress * 105)
                dot_b = int(50 + progress * 100)
                dot_border_color = 'rgb(100,200,100)'
            else:  # red
                dot_r = int(150 + progress * 105)
                dot_g = int(50 + progress * 100)
                dot_b = int(50 + progress * 100)
                dot_border_color = 'rgb(200,100,100)'
            
            dwg.add(dwg.circle(center=(point[1], point[0]), r=4, 
                             fill=f'rgb({dot_r},{dot_g},{dot_b})'))
            dwg.add(dwg.circle(center=(point[1], point[0]), r=6, 
                             fill='none', stroke=dot_border_color, stroke_width=1))
    
    # 绘制起始位置标记 - 蓝色方框 (修正颜色BGR->RGB)
    init_x, init_y = init_pos[1], init_pos[0]
    dwg.add(dwg.rect(insert=(init_x-18, init_y-18), size=(36, 36), 
                   fill='none', stroke='rgb(0,120,180)', stroke_width=6))
    
    # 绘制目标位置 - 多层圆圈效果
    goal_x, goal_y = goal_pos[1], goal_pos[0]
    dwg.add(dwg.circle(center=(goal_x, goal_y), r=24, 
                     fill='none', stroke='rgb(50,200,50)', stroke_width=4))
    dwg.add(dwg.circle(center=(goal_x, goal_y), r=16, fill='rgb(0,180,0)'))
    dwg.add(dwg.circle(center=(goal_x-4, goal_y-4), r=4, fill='rgb(100,255,100)'))
    
    # 可视范围边界指示器已移除 - 与OpenCV界面保持一致
    
    # 绘制agent三角形 - 完整的视觉效果
    agent_x, agent_y = agent_pos[1], agent_pos[0]
    triangle_size = 28
    
    # 根据动作确定三角形朝向
    if agent_action == 0:  # 向右
        points = [(agent_x - triangle_size//2, agent_y - triangle_size//2),
                 (agent_x - triangle_size//2, agent_y + triangle_size//2),
                 (agent_x + triangle_size//2, agent_y)]
        inner_points = [(agent_x - (triangle_size-8)//2, agent_y - (triangle_size-8)//2),
                       (agent_x - (triangle_size-8)//2, agent_y + (triangle_size-8)//2),
                       (agent_x + (triangle_size-8)//2, agent_y)]
    elif agent_action == 1:  # 向左
        points = [(agent_x + triangle_size//2, agent_y - triangle_size//2),
                 (agent_x + triangle_size//2, agent_y + triangle_size//2),
                 (agent_x - triangle_size//2, agent_y)]
        inner_points = [(agent_x + (triangle_size-8)//2, agent_y - (triangle_size-8)//2),
                       (agent_x + (triangle_size-8)//2, agent_y + (triangle_size-8)//2),
                       (agent_x - (triangle_size-8)//2, agent_y)]
    elif agent_action == 2:  # 向下
        points = [(agent_x - triangle_size//2, agent_y - triangle_size//2),
                 (agent_x + triangle_size//2, agent_y - triangle_size//2),
                 (agent_x, agent_y + triangle_size//2)]
        inner_points = [(agent_x - (triangle_size-8)//2, agent_y - (triangle_size-8)//2),
                       (agent_x + (triangle_size-8)//2, agent_y - (triangle_size-8)//2),
                       (agent_x, agent_y + (triangle_size-8)//2)]
    elif agent_action == 3:  # 向上
        points = [(agent_x - triangle_size//2, agent_y + triangle_size//2),
                 (agent_x + triangle_size//2, agent_y + triangle_size//2),
                 (agent_x, agent_y - triangle_size//2)]
        inner_points = [(agent_x - (triangle_size-8)//2, agent_y + (triangle_size-8)//2),
                       (agent_x + (triangle_size-8)//2, agent_y + (triangle_size-8)//2),
                       (agent_x, agent_y - (triangle_size-8)//2)]
    else:  # 默认向上
        points = [(agent_x - triangle_size//2, agent_y + triangle_size//2),
                 (agent_x + triangle_size//2, agent_y + triangle_size//2),
                 (agent_x, agent_y - triangle_size//2)]
        inner_points = [(agent_x - (triangle_size-8)//2, agent_y + (triangle_size-8)//2),
                       (agent_x + (triangle_size-8)//2, agent_y + (triangle_size-8)//2),
                       (agent_x, agent_y - (triangle_size-8)//2)]
    
    # 绘制三角形阴影 - 让阴影比三角形稍大一些
    shadow_size = triangle_size + 4  # 阴影比主体大4px
    if agent_action == 0:  # 向右
        shadow_points = [(agent_x - shadow_size//2 + 4, agent_y - shadow_size//2 + 4),
                        (agent_x - shadow_size//2 + 4, agent_y + shadow_size//2 + 4),
                        (agent_x + shadow_size//2 + 4, agent_y + 4)]
    elif agent_action == 1:  # 向左
        shadow_points = [(agent_x + shadow_size//2 + 4, agent_y - shadow_size//2 + 4),
                        (agent_x + shadow_size//2 + 4, agent_y + shadow_size//2 + 4),
                        (agent_x - shadow_size//2 + 4, agent_y + 4)]
    elif agent_action == 2:  # 向下
        shadow_points = [(agent_x - shadow_size//2 + 4, agent_y - shadow_size//2 + 4),
                        (agent_x + shadow_size//2 + 4, agent_y - shadow_size//2 + 4),
                        (agent_x + 4, agent_y + shadow_size//2 + 4)]
    elif agent_action == 3:  # 向上
        shadow_points = [(agent_x - shadow_size//2 + 4, agent_y + shadow_size//2 + 4),
                        (agent_x + shadow_size//2 + 4, agent_y + shadow_size//2 + 4),
                        (agent_x + 4, agent_y - shadow_size//2 + 4)]
    else:  # 默认向上
        shadow_points = [(agent_x - shadow_size//2 + 4, agent_y + shadow_size//2 + 4),
                        (agent_x + shadow_size//2 + 4, agent_y + shadow_size//2 + 4),
                        (agent_x + 4, agent_y - shadow_size//2 + 4)]
    
    dwg.add(dwg.polygon(points=shadow_points, fill='rgb(30,30,30)'))
    
    # 绘制三角形主体 (修正颜色：OpenCV是BGR，SVG是RGB)
    dwg.add(dwg.polygon(points=points, fill='rgb(255,0,0)', stroke='rgb(150,0,0)', stroke_width=6))
    
    # 绘制内部高光三角形
    dwg.add(dwg.polygon(points=inner_points, fill='rgb(255,100,100)', stroke='rgb(255,150,150)', stroke_width=2))
    
    # 绘制中心亮点
    dwg.add(dwg.circle(center=(agent_x, agent_y), r=4, fill='rgb(255,200,200)'))
    dwg.add(dwg.circle(center=(agent_x, agent_y), r=2, fill='rgb(255,255,255)'))
    
    # ============ 扰动箭头已移除（SVG版本）============
    # 移除扰动箭头功能
    
    # 保存SVG文件
    dwg.save()
    print(f"SVG snapshot saved to {pth}")

def render(grid, state, goal, valid = True, show_vision_mask = False):
        
        state_x = int(state[0])
        state_y = int(state[1])

        food_x = int(goal[0])
        food_y = int(goal[1])

        grid_size_display = 40  # 提升到4倍分辨率 (2x2)
        width, height = grid.shape[0], grid.shape[1]
        img = np.zeros((width * grid_size_display, height * grid_size_display, 3), np.uint8)
        
        for j in range(width):
            for i in range(height):
                if grid[j,i] == 1:
                    # 自由空间 - 改为深色背景
                    cv2.rectangle(img, (i * grid_size_display, j * grid_size_display), 
                                (i * grid_size_display + grid_size_display, j * grid_size_display + grid_size_display), 
                                (30, 40, 50), -1)  # 深蓝灰色
                    
                    # 立体边框效果：深色调
                    # 上边框 - 深色
                    cv2.line(img, (i * grid_size_display, j * grid_size_display), 
                           ((i + 1) * grid_size_display - 1, j * grid_size_display), 
                           (50, 70, 90), 4)
                    # 左边框 - 深色
                    cv2.line(img, (i * grid_size_display, j * grid_size_display), 
                           (i * grid_size_display, (j + 1) * grid_size_display - 1), 
                           (50, 70, 90), 4)
                    # 下边框 - 更深色
                    cv2.line(img, (i * grid_size_display, (j + 1) * grid_size_display - 1), 
                           ((i + 1) * grid_size_display - 1, (j + 1) * grid_size_display - 1), 
                           (15, 25, 35), 4)
                    # 右边框 - 更深色
                    cv2.line(img, ((i + 1) * grid_size_display - 1, j * grid_size_display), 
                           ((i + 1) * grid_size_display - 1, (j + 1) * grid_size_display - 1), 
                           (15, 25, 35), 4)
                else:
                    # 障碍物 - 改为浅色效果
                    cv2.rectangle(img, (i * grid_size_display, j * grid_size_display), 
                                (i * grid_size_display + grid_size_display, j * grid_size_display + grid_size_display), 
                                (200, 200, 210), -1)  # 浅灰色基底
                    
                    # 石墙立体边框效果 (浅色调)
                    cv2.line(img, (i * grid_size_display, j * grid_size_display), 
                           ((i + 1) * grid_size_display - 1, j * grid_size_display), 
                           (230, 230, 240), 4)
                    cv2.line(img, (i * grid_size_display, j * grid_size_display), 
                           (i * grid_size_display, (j + 1) * grid_size_display - 1), 
                           (230, 230, 240), 4)
                    cv2.line(img, (i * grid_size_display, (j + 1) * grid_size_display - 1), 
                           ((i + 1) * grid_size_display - 1, (j + 1) * grid_size_display - 1), 
                           (170, 170, 180), 4)
                    cv2.line(img, ((i + 1) * grid_size_display - 1, j * grid_size_display), 
                           ((i + 1) * grid_size_display - 1, (j + 1) * grid_size_display - 1), 
                           (170, 170, 180), 4)
                    
                    # 添加浅色纹理效果
                    cv2.circle(img, (i * grid_size_display + 12, j * grid_size_display + 12), 2, (180, 180, 190), -1)
                    cv2.circle(img, (i * grid_size_display + 28, j * grid_size_display + 28), 2, (180, 180, 190), -1)
                
                # 智能体位置不在这里绘制，稍后在最上层绘制三角形
                pass
        
        # 目标位置 - 改进的绿色圆形，添加脉冲效果 (放大)
        # 外圈光晕
        cv2.circle(img, (food_y * grid_size_display + grid_size_display//2, food_x * grid_size_display + grid_size_display//2), 
                 24, (50, 200, 50), 4, cv2.LINE_AA)
        # 主体
        cv2.circle(img, (food_y * grid_size_display + grid_size_display//2, food_x * grid_size_display + grid_size_display//2), 
                 16, (0, 180, 0), -1, cv2.LINE_AA)
        # 内部高光
        cv2.circle(img, (food_y * grid_size_display + grid_size_display//2 - 4, food_x * grid_size_display + grid_size_display//2 - 4), 
                 4, (100, 255, 100), -1, cv2.LINE_AA)

        if not valid:
            # draw a big red cross (放大)
            cv2.line(img, (0, 0), (img.shape[1], img.shape[0]), (0, 0, 255), 10, cv2.LINE_AA)
            cv2.line(img, (0, img.shape[0]), (img.shape[1], 0), (0, 0, 255), 10, cv2.LINE_AA)
            cv2.putText(img, "invalid map", (int(img.shape[1]/2) - 200, int(img.shape[0]/2)), cv2.FONT_HERSHEY_SIMPLEX, 2.0, (200, 200, 200), 6, cv2.LINE_AA)
            cv2.putText(img, "invalid map", (int(img.shape[1]/2) - 200, int(img.shape[0]/2)), cv2.FONT_HERSHEY_SIMPLEX, 2.0, (200, 0, 0), 4, cv2.LINE_AA)
        return img

event_type = ""
event_x = 0
event_y = 0

def run_editor(landscape, state, goal, map_size=12):

    global event_type, event_x, event_y

    # define mouse callback function
    def input_cb(event,x,y,flags,param):
        global event_type, event_x, event_y
        if event == cv2.EVENT_MOUSEWHEEL:
            event_x = x
            event_y = y
            event_type = "flip_space"
        elif event == cv2.EVENT_LBUTTONUP:
            event_x = x
            event_y = y
            event_type = "set_start"
        elif event == cv2.EVENT_RBUTTONUP:
            event_x = x
            event_y = y
            event_type = "set_goal"

    cv2.namedWindow("img", cv2.WINDOW_GUI_NORMAL)
    cv2.setMouseCallback("img", input_cb)

    grid = np.array(landscape).reshape(map_size, map_size).transpose()
    valid = True
    grid_size_display = 20

    while True:

        if event_type == "flip_space":
            grid[event_y//grid_size_display, event_x//grid_size_display] = 1 - grid[event_y//grid_size_display, event_x//grid_size_display]
            event_type = "flip_space_done"
        elif event_type == "set_start":
            state = [event_y//grid_size_display, event_x//grid_size_display]
            event_type = "set_start_done"
        elif event_type == "set_goal":
            goal = [event_y//grid_size_display, event_x//grid_size_display]
            event_type = "set_goal_done"

        num_labels, labels, stats, centroids, num_freespace, landscape_img = check_num_labels(grid, map_size, map_size)
        non_zeros = np.count_nonzero(landscape)
        if not (num_labels == 2 and non_zeros >= 5):
            valid = False
        else:
            valid = True
        
        img = render(grid, state, goal, valid)
        cv2.imshow("img", img)
        k = cv2.waitKey(1)
        if k == ord('q'):
            break
        elif k == ord('s'):
            pth = "./logs/test.json"
            grid0 = grid.transpose()
            grid1 = grid0.reshape(map_size*map_size)
            save_current_task([grid1], state[0], state[1], goal[0], goal[1], pth)
            print("task saved to {}".format(pth))
        elif k == ord('r'):
            grid = np.array(landscape).reshape(map_size, map_size).transpose()
    
    grid_ = grid.transpose()
    grid_ = grid_.reshape(map_size*map_size).tolist()
    return grid_


def main():
    """ Manual Control Agent - 键盘控制智能体导航程序
    
    控制说明：
    W/A/S/D - 控制智能体移动（上/左/下/右）
    R - 重置环境和智能体位置
    Q - 退出程序
    E - 进入地图编辑模式
    """

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser(description='Manual Control Agent - 键盘控制智能体导航')
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--path_color", type=str, default="green", choices=["green", "red"], 
                       help="Color scheme for trajectory path rendering")

    args = parser.parse_args()

    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.visualization = args.visualization
    path_color = args.path_color

    """ create landscape
    """
    random_task = True
    # check if file on rpl_config.task_pth exists
    if os.path.isfile(rpl_config.task_pth):
        random_task = False

    if random_task:
        landscape = generate_maze_pool(num_mazes=1, width=10, height=10)
        landscape = padding_landscapes(landscape, width=12, height=12)
    else:
        landscape, state, goal = load_task(pth = rpl_config.task_pth)
        landscape = [landscape]

    print("landscape loaded:")
    print(landscape)

    """ create grid env
    """
    print("Creating grid environment...")
    start_time = time.time()
    GE = GridEnv(landscapes = landscape, width = 12, height = 12, num_envs_per_landscape = 1, reward_free=True)
    GE.reset()
    print("time taken to create envs: ", time.time() - start_time)

    if not random_task:
        # set states of GE
        GE.batched_states = GE.batched_states.at[0, 0].set(state[0])
        GE.batched_states = GE.batched_states.at[0, 1].set(state[1])
        # set goals of GE
        GE.batched_goals = GE.batched_goals.at[0, 0].set(goal[0])
        GE.batched_goals = GE.batched_goals.at[0, 1].set(goal[1])
        GE.init_batched_states, GE.init_batched_goals = jnp.copy(GE.batched_states), jnp.copy(GE.batched_goals)
        GE.batched_goal_reached = batch_compute_goal_reached(GE.batched_states, GE.batched_goals)
        GE.last_batched_goal_reached = jnp.copy(GE.batched_goal_reached)

    render_id = 0
    trajectory = []
    manual_action = 3  # 默认向上
    perturbation_arrows = []  # 存储扰动箭头信息 [(x, y, action, timestamp), ...]
    goal_reached_flag = False  # 标记是否刚到达终点
    svg_saved_flag = False    # 标记是否已保存SVG，避免重复保存
    observation_patches = []  # 记录本轮所有的observation patches

    print("\n=== Manual Control Agent Started ===")
    print("Controls:")
    print("  W/A/S/D - Move agent (Up/Left/Down/Right)")
    print("  R - Reset environment")
    print("  Q - Quit")
    print("  E - Enter map editor")
    print("=====================================\n")
    
    while True:
        """ 键盘控制主循环 """
        
        # 记录当前位置用于轨迹绘制 (适应4倍分辨率)
        trajectory.append([40 * GE.batched_states[render_id][0]+20, 40 * GE.batched_states[render_id][1]+20])

        # 检查是否到达目标
        batched_goal_reached = batch_compute_goal_reached(GE.batched_states, GE.batched_goals)
        
        if batched_goal_reached[render_id]:
            if not svg_saved_flag:  # 只在第一次到达终点时保存SVG
                print("Goal reached! Saving SVG snapshot and patches montage...")
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                
                # 保存主SVG快照
                svg_filename = f"./logs/navigation_snapshot_{timestamp}.svg"
                current_grid = np.array(GE.batched_envs[render_id])
                init_pos = [int(40 * int(GE.init_batched_states[render_id][0]) + 20), 
                           int(40 * int(GE.init_batched_states[render_id][1]) + 20)]
                goal_pos = [int(40 * int(GE.batched_goals[render_id][0]) + 20), 
                           int(40 * int(GE.batched_goals[render_id][1]) + 20)]
                agent_pos = [int(40 * int(GE.batched_states[render_id][0]) + 20), 
                            int(40 * int(GE.batched_states[render_id][1]) + 20)]
                trajectory_native = [[int(point[0]), int(point[1])] for point in trajectory]
                
                save_svg_snapshot(current_grid, trajectory_native, init_pos, goal_pos, agent_pos, manual_action, perturbation_arrows, svg_filename, path_color)
                
                # 保存observation patches拼接图
                if observation_patches:
                    patches_svg_filename = f"./logs/observation_patches_{timestamp}.svg"
                    save_patches_montage_svg(observation_patches, patches_svg_filename)
                    
                    # 创建并显示patches拼接图
                    patches_montage = create_patches_montage(observation_patches)
                    if patches_montage is not None:
                        cv2.imshow("Observation Patches Montage", patches_montage)
                        print(f"Displaying {len(observation_patches)} observation patches")
                
                svg_saved_flag = True  # 标记已保存，避免重复保存
                
            goal_reached_flag = True
            # 不再立即清除轨迹，保持到下一个episode开始
            perturbation_arrows.clear()

        """ 渲染环境 """
        if rpl_config.visualization == "True":
            # 传递当前动作信息给render函数，用于绘制朝向正确的agent三角形
            # 明确设置show_vision_mask=False以去除可视范围显示
            img = GE.render(env_id = render_id, current_action = manual_action, show_vision_mask=False)
            if len(trajectory) > 1:
                # 创建完整的路径序列，确保包含起始位置
                complete_trajectory = []
                
                # 强制插入起始位置（初始位置）
                init_x = int(GE.init_batched_states[render_id][0])
                init_y = int(GE.init_batched_states[render_id][1])
                start_pos = [40 * init_x + 20, 40 * init_y + 20]
                complete_trajectory.append(start_pos)
                
                # 添加所有已记录的轨迹点
                complete_trajectory.extend(trajectory)
                
                # 改进的路径绘制 - 更有质感的效果，支持颜色选择
                for i in range(len(complete_trajectory)-1):
                    # 计算渐变颜色 - 根据path_color参数选择颜色方案
                    progress = i / max(1, len(trajectory) - 2) if len(trajectory) > 2 else 0
                    
                    if path_color == "green":
                        # 渐变绿色：从深绿到亮绿 (OpenCV BGR格式)
                        path_color_bgr = (int(0 + progress * 100),      # B通道
                                         int(100 + progress * 155),     # G通道  
                                         int(0 + progress * 100))       # R通道
                        center_line_color = (150, 255, 150)  # BGR: 绿色
                    else:  # red
                        # 渐变红色：从深红到亮红 (OpenCV BGR格式)
                        path_color_bgr = (int(0 + progress * 100),      # B通道
                                         int(0 + progress * 100),       # G通道  
                                         int(100 + progress * 155))     # R通道
                        center_line_color = (150, 150, 255)  # BGR: 红色
                    
                    # 1. 先绘制阴影效果 (偏移的深色线条)
                    cv2.line(img, 
                            (int(complete_trajectory[i][1]) + 2, int(complete_trajectory[i][0]) + 2), 
                            (int(complete_trajectory[i+1][1]) + 2, int(complete_trajectory[i+1][0]) + 2), 
                            (40, 40, 40), 8, cv2.LINE_AA)
                    
                    # 2. 绘制主路径线条 (较粗的渐变色)
                    cv2.line(img, 
                            (int(complete_trajectory[i][1]), int(complete_trajectory[i][0])), 
                            (int(complete_trajectory[i+1][1]), int(complete_trajectory[i+1][0])), 
                            path_color_bgr, 6, cv2.LINE_AA)
                    
                    # 3. 绘制路径中心的亮线 (更细的亮色线条)
                    cv2.line(img, 
                            (int(complete_trajectory[i][1]), int(complete_trajectory[i][0])), 
                            (int(complete_trajectory[i+1][1]), int(complete_trajectory[i+1][0])), 
                            center_line_color, 2, cv2.LINE_AA)
                
                # 4. 在路径点上绘制小圆点标记
                for i in range(len(complete_trajectory)):
                    progress = i / max(1, len(complete_trajectory) - 1)
                    
                    if path_color == "green":
                        dot_color_bgr = (int(50 + progress * 100),      # B通道
                                        int(150 + progress * 105),      # G通道
                                        int(50 + progress * 100))       # R通道
                        dot_border_color = (100, 200, 100)  # BGR: 绿色
                    else:  # red
                        dot_color_bgr = (int(50 + progress * 100),      # B通道  
                                        int(50 + progress * 100),       # G通道
                                        int(150 + progress * 105))      # R通道
                        dot_border_color = (100, 100, 200)  # BGR: 红色
                    
                    # 绘制路径点的小圆圈
                    cv2.circle(img, (int(complete_trajectory[i][1]), int(complete_trajectory[i][0])), 
                             4, dot_color_bgr, -1, cv2.LINE_AA)
                    cv2.circle(img, (int(complete_trajectory[i][1]), int(complete_trajectory[i][0])), 
                             6, dot_border_color, 1, cv2.LINE_AA)
            
            # ============ 在路径绘制完成后再次绘制Agent三角形，确保不被遮挡 ============
            state_x = int(GE.batched_states[render_id][0])
            state_y = int(GE.batched_states[render_id][1])
            center_x = state_y * 40 + 20  # 40是grid_size_display
            center_y = state_x * 40 + 20
            triangle_size = 28
            
            # 根据动作确定三角形朝向
            if manual_action == 0:  # 向右
                triangle_pts = np.array([
                    [center_x - triangle_size//2, center_y - triangle_size//2],
                    [center_x - triangle_size//2, center_y + triangle_size//2],
                    [center_x + triangle_size//2, center_y]
                ], np.int32)
            elif manual_action == 1:  # 向左
                triangle_pts = np.array([
                    [center_x + triangle_size//2, center_y - triangle_size//2],
                    [center_x + triangle_size//2, center_y + triangle_size//2],
                    [center_x - triangle_size//2, center_y]
                ], np.int32)
            elif manual_action == 2:  # 向下
                triangle_pts = np.array([
                    [center_x - triangle_size//2, center_y - triangle_size//2],
                    [center_x + triangle_size//2, center_y - triangle_size//2],
                    [center_x, center_y + triangle_size//2]
                ], np.int32)
            elif manual_action == 3:  # 向上
                triangle_pts = np.array([
                    [center_x - triangle_size//2, center_y + triangle_size//2],
                    [center_x + triangle_size//2, center_y + triangle_size//2],
                    [center_x, center_y - triangle_size//2]
                ], np.int32)
            else:  # 默认向上
                triangle_pts = np.array([
                    [center_x - triangle_size//2, center_y + triangle_size//2],
                    [center_x + triangle_size//2, center_y + triangle_size//2],
                    [center_x, center_y - triangle_size//2]
                ], np.int32)
            
            # 绘制三角形阴影
            shadow_pts = triangle_pts + 4
            cv2.fillPoly(img, [shadow_pts], (30, 30, 30))
            
            # 绘制三角形外轮廓
            cv2.polylines(img, [triangle_pts], True, (0, 0, 150), 6, cv2.LINE_AA)
            
            # 绘制三角形主体
            cv2.fillPoly(img, [triangle_pts], (0, 0, 255))
            
            # 绘制内部高光
            inner_size = triangle_size - 8
            if manual_action == 0:  # 向右
                inner_triangle_pts = np.array([
                    [center_x - inner_size//2, center_y - inner_size//2],
                    [center_x - inner_size//2, center_y + inner_size//2],
                    [center_x + inner_size//2, center_y]
                ], np.int32)
            elif manual_action == 1:  # 向左
                inner_triangle_pts = np.array([
                    [center_x + inner_size//2, center_y - inner_size//2],
                    [center_x + inner_size//2, center_y + inner_size//2],
                    [center_x - inner_size//2, center_y]
                ], np.int32)
            elif manual_action == 2:  # 向下
                inner_triangle_pts = np.array([
                    [center_x - inner_size//2, center_y - inner_size//2],
                    [center_x + inner_size//2, center_y - inner_size//2],
                    [center_x, center_y + inner_size//2]
                ], np.int32)
            elif manual_action == 3:  # 向上
                inner_triangle_pts = np.array([
                    [center_x - inner_size//2, center_y + inner_size//2],
                    [center_x + inner_size//2, center_y + inner_size//2],
                    [center_x, center_y - inner_size//2]
                ], np.int32)
            else:  # 默认向上
                inner_triangle_pts = np.array([
                    [center_x - inner_size//2, center_y + inner_size//2],
                    [center_x + inner_size//2, center_y + inner_size//2],
                    [center_x, center_y - inner_size//2]
                ], np.int32)
            
            cv2.fillPoly(img, [inner_triangle_pts], (100, 100, 255))
            cv2.polylines(img, [inner_triangle_pts], True, (150, 150, 255), 2, cv2.LINE_AA)
            cv2.circle(img, (center_x, center_y), 4, (200, 200, 255), -1, cv2.LINE_AA)
            cv2.circle(img, (center_x, center_y), 2, (255, 255, 255), -1, cv2.LINE_AA)

            cv2.imshow("Manual Control Agent", img)

            # 键盘控制处理
            k = cv2.waitKey(1) & 0xFF
            if k == ord('r'): 
                print("Resetting environment...")
                trajectory.clear()
                observation_patches.clear()  # 清除observation patches
                goal_reached_flag = False
                svg_saved_flag = False  # 重置SVG保存标志
                GE.reset()
                if not random_task:
                    # set states of GE
                    GE.batched_states = GE.batched_states.at[0, 0].set(state[0])
                    GE.batched_states = GE.batched_states.at[0, 1].set(state[1])
                    # set goals of GE
                    GE.batched_goals = GE.batched_goals.at[0, 0].set(goal[0])
                    GE.batched_goals = GE.batched_goals.at[0, 1].set(goal[1])
                    GE.init_batched_states, GE.init_batched_goals = jnp.copy(GE.batched_states), jnp.copy(GE.batched_goals)
                    GE.batched_goal_reached = batch_compute_goal_reached(GE.batched_states, GE.batched_goals)
                    GE.last_batched_goal_reached = jnp.copy(GE.batched_goal_reached)

            elif k == ord('q'):
                print("Exiting...")
                break
                
            elif k == ord('w'):  # 向上移动
                # 如果刚到达终点，先清除历史数据开始新episode
                if goal_reached_flag:
                    trajectory.clear()
                    observation_patches.clear()  # 清除observation patches
                    goal_reached_flag = False
                    svg_saved_flag = False  # 重置SVG保存标志
                    
                manual_action = 3
                batched_actions = jnp.array([manual_action])
                batched_goal_reached, _ = GE.step(batched_actions, reset=True)
                
                # 在执行动作后记录新的observation
                if hasattr(GE, 'concat_obs') and GE.concat_obs is not None:
                    current_obs = np.array(GE.concat_obs[render_id][:9])  # 只取前9个元素(3x3观测)
                    observation_patches.append(current_obs.tolist())
                    
                print("Moving UP")
                
            elif k == ord('s'):  # 向下移动
                # 如果刚到达终点，先清除历史数据开始新episode
                if goal_reached_flag:
                    trajectory.clear()
                    observation_patches.clear()  # 清除observation patches
                    goal_reached_flag = False
                    svg_saved_flag = False  # 重置SVG保存标志
                    
                manual_action = 2
                batched_actions = jnp.array([manual_action])
                batched_goal_reached, _ = GE.step(batched_actions, reset=True)
                
                # 在执行动作后记录新的observation
                if hasattr(GE, 'concat_obs') and GE.concat_obs is not None:
                    current_obs = np.array(GE.concat_obs[render_id][:9])  # 只取前9个元素(3x3观测)
                    observation_patches.append(current_obs.tolist())
                    
                print("Moving DOWN")
                
            elif k == ord('a'):  # 向左移动
                # 如果刚到达终点，先清除历史数据开始新episode
                if goal_reached_flag:
                    trajectory.clear()
                    observation_patches.clear()  # 清除observation patches
                    goal_reached_flag = False
                    svg_saved_flag = False  # 重置SVG保存标志
                    
                manual_action = 1
                batched_actions = jnp.array([manual_action])
                batched_goal_reached, _ = GE.step(batched_actions, reset=True)
                
                # 在执行动作后记录新的observation
                if hasattr(GE, 'concat_obs') and GE.concat_obs is not None:
                    current_obs = np.array(GE.concat_obs[render_id][:9])  # 只取前9个元素(3x3观测)
                    observation_patches.append(current_obs.tolist())
                    
                print("Moving LEFT")
                
            elif k == ord('d'):  # 向右移动
                # 如果刚到达终点，先清除历史数据开始新episode
                if goal_reached_flag:
                    trajectory.clear()
                    observation_patches.clear()  # 清除observation patches
                    goal_reached_flag = False
                    svg_saved_flag = False  # 重置SVG保存标志
                    
                manual_action = 0
                batched_actions = jnp.array([manual_action])
                batched_goal_reached, _ = GE.step(batched_actions, reset=True)
                
                # 在执行动作后记录新的observation
                if hasattr(GE, 'concat_obs') and GE.concat_obs is not None:
                    current_obs = np.array(GE.concat_obs[render_id][:9])  # 只取前9个元素(3x3观测)
                    observation_patches.append(current_obs.tolist())
                    
                print("Moving RIGHT")
                
            elif k == ord('e'):
                print("Entering map editor...")
                new_landscape = run_editor(landscape[0], GE.batched_states[0], GE.batched_goals[0])
                GE.set_landscapes_only([new_landscape])
                landscape[0] = new_landscape

    cv2.destroyAllWindows()
    print("Manual Control Agent terminated.")

if __name__ == "__main__":
    main()
